# 数据增强效果
# %cd PaddleClas/
from ppcls.data.imaug import ToCHWImage
from ppcls.data.imaug import DecodeImage
from ppcls.data.imaug import ResizeImage
from ppcls.data.imaug import RandAugment
from ppcls.data.imaug import transform
from ppcls.data.imaug import MixupOperator
from ppcls.data.imaug import CutmixOperator
import numpy as np
import os
from matplotlib import pyplot as plt
from PIL import Image
import cv2

size = 224
imgs = []
decode_op = DecodeImage()
resize_op = ResizeImage(size=(size, size))
tochw_op = ToCHWImage()
mixup_op = MixupOperator()
imgs_dir = "./dataset/fmd/glass/"
save_path_1 = "./dataAugm/CutMix"
fnames = os.listdir(imgs_dir)
mixup_op = MixupOperator()
cutmix_op = CutmixOperator()
ops = [decode_op, resize_op, tochw_op]
batch = []

for idx, f in enumerate(fnames):
    data = open(os.path.join(imgs_dir, f),"rb").read()
    img = transform(data, ops)    
    
new_batch = cutmix_op(batch)

for bt in new_batch:
    img = bt[0]
    save_path = os.path.join(save_path_1,str(bt[1])+".jpg")
    print(save_path)
    img = img.transpose((1, 2, 0))
    cv2.imwrite(save_path, img)
